import torch
torch.set_num_threads(1)
from torch.utils.data import Dataset
import pandas as pd
from sklearn.preprocessing import StandardScaler, LabelEncoder

# 预处理函数
def preprocess_data(df, label="cogimp_label", fit=True, scaler=None, cat_encoders=None):
    """
    :param df: 数据集
    :param label: 标签列
    :param fit: 是否拟合scaler和cat_encoders
    :param scaler: 标准化器
    :param cat_encoders: 类别特征编码器
    :return: 处理后的数据集，标准化器，类别特征编码器，连续特征列，类别特征列
    """
    df = df.copy()
    categorical_cols = df.select_dtypes(include=['object', 'category']).columns
    numerical_cols = df.select_dtypes(include=['int64', 'float64']).columns
    numerical_cols = [c for c in numerical_cols if c != label]

    # 类别特征编码
    if fit:
        cat_encoders = {}
        for col in categorical_cols:
            le = LabelEncoder()
            df[col] = le.fit_transform(df[col].astype(str))
            cat_encoders[col] = le
    else:
        for col in categorical_cols:
            le = cat_encoders[col]
            df[col] = le.transform(df[col].astype(str))

    # 连续特征标准化
    if fit:
        scaler = StandardScaler()
        df[numerical_cols] = scaler.fit_transform(df[numerical_cols])
    else:
        df[numerical_cols] = scaler.transform(df[numerical_cols])

    return df, scaler, cat_encoders, numerical_cols, categorical_cols


# Dataset
class ADHDDataset(Dataset):
    def __init__(self, df, label="cogimp_label", task="classification", fit=True, scaler=None, cat_encoders=None):
        """
        :param df: 输入 DataFrame
        :param label: 标签列
        :param task: "classification" or "regression"
        """
        self.task = task
        self.df, self.scaler, self.cat_encoders, self.num_cols, self.cat_cols = preprocess_data(
            df, label=label, fit=fit, scaler=scaler, cat_encoders=cat_encoders
        )

        # 根据任务设置标签类型
        if self.task == "classification":
            self.labels = torch.tensor(self.df[label].values, dtype=torch.long)
        else:
            self.labels = torch.tensor(self.df[label].values, dtype=torch.float)

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # 类别特征 → LongTensor（embedding用）
        x_cat = torch.tensor(self.df.iloc[idx][self.cat_cols].values, dtype=torch.long)
        # 连续特征 → FloatTensor
        x_cont = torch.tensor(self.df.iloc[idx][self.num_cols].values, dtype=torch.float)
        # 标签
        y = self.labels[idx]
        return x_cat, x_cont, y
